package edu.gatech.cc.liam.core.rl.models;

import java.util.Arrays;
import java.util.Iterator;
import java.util.NoSuchElementException;

public class CartesianIterator implements Iterator<int[]>, Iterable<int[]> {

	private int[] current;
	int[] setSizes;
	int numElementsReturn;
	int numElements;
	
	public CartesianIterator(int[] setSizes) {
		super();
		this.setSizes = setSizes;
		current = new int[setSizes.length];
		Arrays.fill(current, 0);
		//the index gets incremented during the first call to next()
		current[0] = -1;
		numElementsReturn = 0;
		numElements = 1;
		for(int i=0; i<setSizes.length; i++)
			numElements *= setSizes[i];
	}

	public static int getJointIndex(int[] fieldSizes, int[] indices) {
		int indexVal = indices[0];
		int multiplier = 1;
		for(int i = 1; i < indices.length; i++) {
			assert(indices[i] < fieldSizes[i]);
			multiplier *= fieldSizes[i-1];
			indexVal += indices[i] * multiplier;
		}
		return indexVal;
	}
	
	@Override
	public boolean hasNext() {
		return numElementsReturn < numElements;
	}

	@Override
	public int[] next() {
		current[0] ++;
		// Carry if needed
		for(int i=0; i<current.length-1; i++) {
			if(current[i] == setSizes[i]) {
				current[i]=0;
				current[i+1]++;
			}
		}
		if(current[current.length-1] == setSizes[current.length-1]) {
			throw new NoSuchElementException();
		}
		numElementsReturn++;
		return current.clone();
	}

	public int[] getCurrent() {
		return current.clone();
	}
	
	public int getElementsReturned() {
		return numElementsReturn;
	}
	
	@Override
	public void remove() {
		throw new UnsupportedOperationException();			
	}

	@Override
	public Iterator<int[]> iterator() {
		return this;
	}
	
}
